###########################################################
#     Prj: WGCNA all in one
#     Assignment: check scale free
#     Author: Shawn Wang
#     Date: Mar 24,2021
###########################################################

# options and package -----------------------------------------------------
suppressMessages(library(ggplot2))
suppressMessages(library(stringr))
suppressMessages(library(reshape2))
suppressMessages(library(ggprism))
suppressMessages(library(patchwork))
suppressMessages(library(tidyverse))
suppressMessages(library(WGCNA))
suppressMessages(library(getopt))
allowWGCNAThreads()
options(stringsAsFactors = F)
# args --------------------------------------------------------------------

command=matrix(c(
  'help', 'h', 0, 'logic', 'help information',
  'powerIn', 'p', 1, 'integer', 'test power by last step',
  'minModuleSize', 'm', 1, 'integer', 'Minimum number of genes in the module',
  'mergeCutHeight', 'c', 1, 'double', 'Dynamic cutting height',
  'traitData', 't', 1, 'character', 'trait data'
),byrow = T, ncol = 5)
args = getopt(command)

## default valuels==========
if (!is.null(args$help)) {
  cat(paste(getopt(command, usage = T), "\n"))
  #  q(status=1)
}

powerIn = args$powerIn

## default value
if (is.null(args$minModuleSize)){
  args$minModuleSize = 30
}
minModuleSize = args$minModuleSize

if (is.null(args$mergeCutHeight)){
  args$mergeCutHeight = 0.25
}
if (is.null(args$powerIn)){
  q(status=1)
}
mergeCutHeight = args$mergeCutHeight 

traitData = args$traitData 


# out message -------------------------------------------------------------
con <- file("WGCNA.AllInOneStep2.log") 
sink(con, append=TRUE) 
sink(con, append=TRUE, type="message") 
## test
load("WGCNAoneInAll.Rdata")
# load("/Volumes/Samsung_T5/毕业论文/04.WGCNA/01.rawdata/test/WGCNAoneInAll.Rdata")
load("~/02.MyScript/OneStepWGCNA/04.OneInAll/functions.Rdata")
nGenes = ncol(datExpr)
nSample = nrow(datExpr)
# ## test
# powerIn = 14
# mergeCutHeight = 0.4
# minModuleSize = 100
# traitData = "/Volumes/Samsung_T5/毕业论文/04.WGCNA/01.rawdata/fiber.yandq.trait.txt"
# net para ----------------------------------------------------------------
para1 = getnetwork(datExpr = datExpr,power = powerIn,minModuleSize = minModuleSize,mergeCutHeight = mergeCutHeight)
## save para=======
net = para1$net
moduleLabels = para1$moduleLabels
moduleColors = para1$moduleColors
MEs_col = para1$MEs_col
MEs = para1$MEs
Gene2module = para1$Gene2module
write.table(Gene2module,file = "02.gene2module.xls",row.names = F,sep = "\t",quote = F)
## save Dendrograms plot=========
pdf(file = "02.NetDendrograms.pdf",width = 8,height = 7)
plotDendroAndColors(net$dendrograms[[1]], moduleColors[net$blockGenes[[1]]],
                    "Module colors",
                    dendroLabels = FALSE, hang = 0.03,
                    addGuide = TRUE, guideHang = 0.05)
dev.off()

write.table(table(moduleColors),file = "02.moduleColor.txt",row.names = F,sep = "\t") 
## save adjacency plot==========
pdf(file = "02.EigengeneAdjHeatmap.pdf",width = 7,height = 10)
plotEigengeneNetworks(MEs_col, "Eigengene adjacency heatmap", 
                      marDendro = c(3,3,2,4),
                      marHeatmap = c(3,4,2,2), plotDendrograms = T, 
                      xLabelsAngle = 90)
dev.off()
# module Trait ------------------------------------------------------------
trait = read.table(traitData,header = T,sep = "\t")
if (ncol(trait) == 2) {
  x <- trait
  Tcol = as.character(unique(x[,2]))
  b <- list()
  for (i in 1:length(Tcol)) {
    b[[i]] = data.frame(row.names = x[,1],
                        levels = ifelse(x[,2] == Tcol[i],1,0))
  }
  c <- bind_cols(b)
  c <- data.frame(row.names = x$name,
                  c)
  colnames(c) = Tcol
  rownames(c) = trait[,1]
  pheTmp <- c
} else {
  pheTmp = data.frame(row.names = trait[,1],
                      trait[,-1])
}
pheTmp = pheTmp[match(rownames(datExpr),rownames(pheTmp)),]

Mt.Tmp = getMt(phenotype = pheTmp,MEs_col = MEs_col,nSamples = nSamples,moduleColors = moduleColors,datExpr = datExpr)

modTraitCor = Mt.Tmp$modTraitCor
modTraitP = Mt.Tmp$modTraitP
textMatrix = Mt.Tmp$textMatrix
pdf(file = "03.Module-traitRelation.pdf",width = 10,height = 10)
labeledHeatmap(Matrix = modTraitCor, xLabels = colnames(pheTmp), 
               yLabels = colnames(MEs_col), 
               cex.lab = 0.7, xLabelsAngle = 45, xLabelsAdj = 1,
               ySymbols = substr(colnames(MEs_col),3,1000), colorLabels = FALSE, 
               colors = colorRampPalette(c("orange","white","purple"))(100), 
               textMatrix = textMatrix, setStdMargins = FALSE, 
               cex.text = 0.6, zlim = c(-1,1),
               main = paste("Module-trait relationships"))
dev.off()
## KME and MM
datKME = getKME(datExpr = datExpr,moduleColors = moduleColors,MEs_col = MEs_col)
MM.para = getMM(datExpr = datExpr,MEs_col = MEs_col,nSamples = nSample,corType = corType)
MM = MM.para$MM
MMP = MM.para$MMP

write.table(datKME,file = "02.KMEofAllModule.xls",quote = F,sep = "\t",row.names = T)

## all verbosScatterPlot
traitName = colnames(pheTmp)
plist = list()
for (i in 1:length(traitName)) {
  plist[[i]] = MMvsGSall(which.trait = traitName[i],traitData = pheTmp,datExpr = datExpr,
                         moduleColors = moduleColors,geneModuleMembership = MM,MEs = MEs_col,nSamples = nSamples)
  ggsave(plist[[i]],filename = paste0("04.",traitName[i],".GSvsMM.pdf"),width = 25,height = as.integer(0.8+length(traitName)/5)*10)
}


# hubgene -----------------------------------------------------------------
hublist = list()
eachTrait = list()

colorlevels = unique(moduleColors)
for (i in 1:ncol(pheTmp)) {
  for (j in 1:length(colorlevels)) {
    eachTrait[[j]] = hubgenes(datExpr = datExpr,mdl = colorlevels[j],power = powerIn,trt =colnames(pheTmp)[i] ,KME = datKME,GS.cut = 0.5,kME.cut = 0.5,datTrait = pheTmp)
    eachTrait[[j]]$module = colorlevels[j]
    eachTrait[[j]]$trait = colnames(pheTmp)[i]
  }
  hublist[[i]] = bind_rows(eachTrait)
  eachTrait = list()
}

hubout = bind_rows(hublist)

write.table(hubout,"05.hubgene.xls",row.names = F,sep = "\t",quote = F)

save.image("step2.Rdata")


sink()
sink(type="message")

cat(readLines("WGCNA.AllInOneStep2.log"), sep="\n")
write.table(cat(readLines("WGCNA.AllInOneStep2.log"), sep="\n"), "log.txt")
